import pymysql
import os
from werkzeug.utils import secure_filename

class MerchantManager:
    def __init__(self):
        self.mysql_pwd = os.getenv('MYSQL_PWD', "123456")
        self.db_name = os.getenv('DB_NAME', "appDB")
        self.upload_folder = '/static/images/'
        self.allowed_extensions = set(['png', 'jpg', 'jpeg'])

    def get_db_connection(self):
        """获取数据库连接"""
        return pymysql.connect(
            host="localhost",
            user="root",
            password=self.mysql_pwd,
            database=self.db_name,
            charset='utf8'
        )

    def allowed_file(self, filename):
        """检查文件扩展名是否允许"""
        return '.' in filename and filename.rsplit('.', 1)[1].lower() in self.allowed_extensions

    def get_menu_list(self, username):
        """获取餐厅菜单列表"""
        db = None
        try:
            db = self.get_db_connection()
            cursor = db.cursor()
            sql = "SELECT * FROM DISHES WHERE restaurant = %s"
            cursor.execute(sql, (username,))
            res = cursor.fetchall()
            return res if res else None
        except Exception as e:
            print(f"Error: {e}")
            return None
        finally:
            if db:
                db.close()

    def delete_dish(self, dishname, restaurant):
        """删除菜品"""
        db = None
        try:
            db = self.get_db_connection()
            cursor = db.cursor()
            sql = "DELETE FROM DISHES WHERE dishname = %s AND restaurant = %s"
            cursor.execute(sql, (dishname, restaurant))
            db.commit()
            return True
        except Exception as e:
            print(f"Error: {e}")
            return False
        finally:
            if db:
                db.close()

    def sort_menu_by_sales(self, username):
        """按销量排序菜单"""
        db = None
        try:
            db = self.get_db_connection()
            cursor = db.cursor()
            sql = "SELECT * FROM DISHES WHERE restaurant = %s ORDER BY sales DESC"
            cursor.execute(sql, (username,))
            res = cursor.fetchall()
            return res if res else None
        except Exception as e:
            print(f"Error: {e}")
            return None
        finally:
            if db:
                db.close()

    def sort_menu_by_price(self, username):
        """按价格排序菜单"""
        db = None
        try:
            db = self.get_db_connection()
            cursor = db.cursor()
            sql = "SELECT * FROM DISHES WHERE restaurant = %s ORDER BY price DESC"
            cursor.execute(sql, (username,))
            res = cursor.fetchall()
            return res if res else None
        except Exception as e:
            print(f"Error: {e}")
            return None
        finally:
            if db:
                db.close()

    def update_dish(self, dishname, restaurant, dishinfo, nutriention, price, isSpecialty, imgsrc=None):
        """更新菜品信息"""
        db = None
        try:
            db = self.get_db_connection()
            cursor = db.cursor()
            if imgsrc:
                sql = """
                    UPDATE DISHES 
                    SET dishinfo = %s, nutriention = %s, price = %s, imgsrc = %s, isSpecialty = %s 
                    WHERE dishname = %s AND restaurant = %s
                """
                cursor.execute(sql, (dishinfo, nutriention, price, imgsrc, isSpecialty, dishname, restaurant))
            else:
                sql = """
                    UPDATE DISHES 
                    SET dishinfo = %s, nutriention = %s, price = %s, isSpecialty = %s 
                    WHERE dishname = %s AND restaurant = %s
                """
                cursor.execute(sql, (dishinfo, nutriention, price, isSpecialty, dishname, restaurant))
            db.commit()
            return True
        except Exception as e:
            print(f"Error: {e}")
            return False
        finally:
            if db:
                db.close()

    def add_dish(self, dishname, restaurant, dishinfo, nutriention, price, imgsrc, isSpecialty):
        """添加菜品"""
        db = None
        try:
            db = self.get_db_connection()
            cursor = db.cursor()
            sql = """
                INSERT INTO DISHES (dishname, restaurant, dishinfo, nutriention, price, imgsrc, isSpecialty) 
                VALUES (%s, %s, %s, %s, %s, %s, %s)
            """
            cursor.execute(sql, (dishname, restaurant, dishinfo, nutriention, price, imgsrc, isSpecialty))
            db.commit()
            return True
        except Exception as e:
            print(f"Error: {e}")
            return False
        finally:
            if db:
                db.close()